{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%% Load data\n"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from skimage.io import imread\n",
    "from network.wrapper import GIFTDescriptor\n",
    "from train.evaluation import EvaluationWrapper, Matcher\n",
    "from utils.superpoint_utils import SuperPointWrapper, SuperPointDescriptor\n",
    "from utils.base_utils import draw_correspondence, draw_keypoints\n",
    "from utils.match_utils import compute_angle\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "\n",
    "st_peters_idxs=[12,26,30,42,98]\n",
    "pair_id=4 # you can try 0-4\n",
    "\n",
    "M=np.loadtxt(f'demo/st_peters_{st_peters_idxs[pair_id]}.txt')\n",
    "K0, K1, R = M[:3,:], M[3:6,:], M[6:,:]\n",
    "img0=imread(f'demo/st_peters_{st_peters_idxs[pair_id]}_left.jpg')\n",
    "img1=imread(f'demo/st_peters_{st_peters_idxs[pair_id]}_right.jpg')\n",
    "\n",
    "def show_img(img0,img1):\n",
    "    h0,h1=img0.shape[0],img1.shape[0]\n",
    "    if h0<h1: img0=cv2.copyMakeBorder(img0,0,h1-h0,0,0,borderType=cv2.BORDER_CONSTANT,value=0)\n",
    "    if h1<h0: img1=cv2.copyMakeBorder(img1,0,h0-h1,0,0,borderType=cv2.BORDER_CONSTANT,value=0)\n",
    "    img=np.concatenate([img0,img1],axis=1)\n",
    "    plt.figure(figsize=(20,16))\n",
    "    plt.imshow(img)\n",
    "    plt.show()\n",
    "\n",
    "show_img(img0,img1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%% \n"
    }
   },
   "outputs": [],
   "source": [
    "# Load detector and descriptor\n",
    "\n",
    "detector=SuperPointWrapper(EvaluationWrapper.load_cfg('configs/eval/superpoint_det.yaml'))\n",
    "gift_desc=GIFTDescriptor(EvaluationWrapper.load_cfg('configs/eval/gift_pretrain_desc.yaml'))\n",
    "superpoint_desc=SuperPointDescriptor(EvaluationWrapper.load_cfg('configs/eval/superpoint_desc.yaml'))\n",
    "matcher=Matcher(EvaluationWrapper.load_cfg('configs/eval/match_v0.yaml'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%% \n"
    }
   },
   "outputs": [],
   "source": [
    "# Detection\n",
    "\n",
    "kps0,_=detector(img0)\n",
    "kps1,_=detector(img1)\n",
    "kps_img0=draw_keypoints(img0, kps0, radius=4)\n",
    "kps_img1=draw_keypoints(img1, kps1, radius=4)\n",
    "\n",
    "h0,h1=kps_img0.shape[0],kps_img1.shape[0]\n",
    "if h0<h1: kps_img0=cv2.copyMakeBorder(kps_img0,0,h1-h0,0,0,borderType=cv2.BORDER_CONSTANT,value=0)\n",
    "if h1<h0: kps_img1=cv2.copyMakeBorder(kps_img1,0,h0-h1,0,0,borderType=cv2.BORDER_CONSTANT,value=0)\n",
    "kps_img=np.concatenate([kps_img0,kps_img1],axis=1)\n",
    "plt.figure(figsize=(20,16))\n",
    "plt.imshow(kps_img)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%% \n"
    }
   },
   "outputs": [],
   "source": [
    "# relative rotation estimation SuperPoint Desc\n",
    "\n",
    "sup_desc0=superpoint_desc(img0,kps0)\n",
    "sup_desc1=superpoint_desc(img1,kps1)\n",
    "sup_matches=matcher.match(sup_desc0,sup_desc1)\n",
    "sup_inlier_mask, sup_R, _ =EvaluationWrapper.estimate_relative_pose_from_correspondence(kps0[sup_matches[:,0]],kps1[sup_matches[:,1]],K0,K1)\n",
    "sup_diff_angle=compute_angle(sup_R.T @ R)\n",
    "print(f'superpoint: the angle between ground truth rotation and the estimated rotation is {sup_diff_angle:.2f} degree')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%% \n"
    }
   },
   "outputs": [],
   "source": [
    "# relative rotation estimation SuperPoint Desc\n",
    "\n",
    "gift_desc0=gift_desc(img0,kps0)\n",
    "gift_desc1=gift_desc(img1,kps1)\n",
    "gift_matches=matcher.match(gift_desc0,gift_desc1)\n",
    "gift_inlier_mask, gift_R, _ =EvaluationWrapper.estimate_relative_pose_from_correspondence(kps0[gift_matches[:,0]],kps1[gift_matches[:,1]],K0,K1)\n",
    "gift_diff_angle=compute_angle(gift_R.T @ R)\n",
    "print(f'GIFT: the angle between ground truth rotation and the estimated rotation is {gift_diff_angle:.2f} degree')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# draw inlier correspondences\n",
    "gift_img=draw_correspondence(img0,img1,kps0,kps1,gift_matches[gift_inlier_mask])\n",
    "sup_img=draw_correspondence(img0,img1,kps0,kps1,sup_matches[sup_inlier_mask])\n",
    "\n",
    "plt.figure(figsize=(40,32))\n",
    "plt.subplot(211); plt.imshow(gift_img); plt.title(f'GIFT error {gift_diff_angle:.2f} degree')\n",
    "plt.subplot(212); plt.imshow(sup_img); plt.title(f'superpoint error {sup_diff_angle:.2f} degree')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
